from segment_anything import SamPredictor, sam_model_registry
import cv2  # For loading the image
import numpy as np
import time


def load_model(checkpoint_path = '', model_type = "vit_h"):
    pass
    sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
    predictor = SamPredictor(sam)

    return predictor


def do_infer(input_points, input_labels, predictor, image_path):
    start_time = time.time()
    image = cv2.imread(image_path)
    predictor.set_image(image)
    input_points = np.array(input_points)  # 使用 np.array() 方法進行轉換
    # input_labels = np.array(input_labels)

    masks, scores, logits = predictor.predict(
        point_coords=input_points,
        point_labels=input_labels,
        multimask_output=False  # Set to True if you want multiple masks
    )   

    # Output the resulting masks and scores
    # print("Generated Masks:", masks)
    # print("Confidence Scores:", scores)
    end_time = time.time()
    elapsed_time = end_time - start_time  # 计算耗时

    print(f"do_infer 耗时: {elapsed_time:.4f} 秒")  # 打印耗时

    return masks, scores, logits


if __name__ == '__main__':

# Define input prompts (e.g., points and their labels)
    input_points = [[774, 308]]  # Example: Two points in the image
    input_labels = [1]  # 1 for foreground, 0 for background

    checkpoint_path = '/home/JSDC/017254/code/gitee/deep-learing/model/sam_vit_h_4b8939.pth'
    image_path = "/home/JSDC/017254/code/gitee/deep-learing/imgs/1744185603703480.jpg"  # Replace with your actual image path
    predictor = load_model(checkpoint_path)
    
    masks, scores, logits = do_infer(input_points, input_labels, predictor, image_path)

    for i in range(masks.shape[0]):
        mask = masks[i]  # 獲取每個遮罩 (形狀為 768x1024)
        
        # 將遮罩數據轉換為 0-255 的範圍（如果不是）
        mask = (mask * 255).astype(np.uint8)  # 確保數據是整數類型，範圍在 0-255
        
        # 保存遮罩為圖片
        output_filename = f"masks/mask_{i}.png"  # 生成文件名
        cv2.imwrite(output_filename, mask)  # 使用 OpenCV 保存圖片
        print(f"Saved mask {i} as {output_filename}")
    pass
